# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/19 15:24 
@Author  : ZhangShenao 
@File    : 1.数据导入与相似性检索.py 
@Desc    : 使用LangChain封装的Faiss向量数据库,实现数据导入与相似性检索
"""
import dotenv
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.vectorstores import FAISS

# 加载环境变量
dotenv.load_dotenv()

# 创建ZhipuAIEmbeddings
# 使用embedding-2这个Embedding模型
embeddings = ZhipuAIEmbeddings(model='embedding-2')

# 从文本导入Faiss向量数据库
faiss = FAISS.from_texts(texts=[
    '笨笨是一只很喜欢睡觉的猫咪',
    '我喜欢在夜晚听音乐，这让我感到放松。',
    '猫咪在窗台上打盹，看起来非常可爱。',
    '学习新技能是每个人都应该追求的目标。',
    '我最喜欢的食物是意大利面，尤其是番茄酱的那种。',
    '昨晚我做了一个奇怪的梦，梦见自己在太空飞行。',
    '我的手机突然关机了，让我有些焦虑。',
    '阅读是我每天都会做的事情，我觉得很充实。',
    '他们一起计划了一次周末的野餐，希望天气能好。',
    '我的狗喜欢追逐球，看起来非常开心。',
], embedding=embeddings)

# 打印索引下的向量数量
# print(faiss.index.ntotal)  # 10

# 进行基础的相似度检索
similarity_search_result = faiss.similarity_search(
    query='我养了一只猫，它叫笨笨',  # 指定检索关键词
    k=2,  # 指定返回的数据条数
)
print(f'similarity_search_result: {similarity_search_result}')

# 带有得分的相似度检索,返回的得分是欧氏距离,距离越短表示越相关
similarity_search_with_score_result = faiss.similarity_search_with_score(
    query='我养了一只猫，它叫笨笨',  # 指定检索关键词
    k=2,  # 指定返回的数据条数
)
print(f'similarity_search_with_score_result: {similarity_search_with_score_result}')

# 带有相关性得分的相似度检索,返回的是相似性得分
# 得分在[0,1]范围内
# 相似性得分越高,表示越相关
similarity_search_with_relevance_score_result = faiss.similarity_search_with_relevance_scores(
    query='我养了一只猫，它叫笨笨',  # 指定检索关键词
    k=5,  # 指定返回的数据条数
)
print(f'similarity_search_with_relevance_score_result: {similarity_search_with_relevance_score_result}')
